# coding=utf-8
import os
import shutil
import sys
import threading
from queue import Queue

import pandas as pd
import akshare as ak

from loguru import logger

from cal_ops.point import cal_point1_txt
from models.stock_model import StockNumber
from mylib import download_all

st_queue = Queue()


def get_all_txt_files(k_word):
    res = []
    for root, dirs, files in os.walk('txts'):
        res = [(os.path.join(root, item), item.replace('.txt', '')) for item in files]
    if k_word:
        return [item for item in res if k_word in item]
    return res


def clear_dir(d='result'):
    if os.path.exists(d):
        shutil.rmtree(d)
    os.makedirs(d)


def get_all_stock_csv_path():
    for root, dirs, files in os.walk('stocks'):
        return [os.path.join(root, item) for item in files]


def get_all_stocks_indus():
    with open('industry_count_rank.csv', 'r') as fr:
        all_indus = [item.strip().split(',') for item in fr.readlines()[1:]]
        return all_indus


def get_avg_top10(stocks, N=50):
    stocks_top_dict = {}
    for stock in stocks:
        csv_path = f'stocks/{stock}.csv'
        if not os.path.exists(csv_path):
            continue
        df2 = pd.read_csv(csv_path)
        stocks_top_dict[stock] = round(abs(df2[0:N]['pct_chg'].mean()), 4)
    sorted_d = sorted(stocks_top_dict.items(), key=lambda x: x[1], reverse=True)
    # 获取近N日振幅排名前10的股票
    return sorted_d[:10]


def get_end_date(log_dir, indus, indus_num):
    f_path = f'{log_dir}/{indus}_all_{indus_num}.csv'
    if not os.path.exists(f_path):
        return None
    with open(f_path, 'r') as fr:
        line = fr.readline()
        return eval(line.split(',')[0])


def t_download(today_date, kw):
    txts = get_all_txt_files(kw)
    for txt_path, txt_name in txts:
        with open(txt_path) as fr:
            stocks_names = [item.strip() for item in fr.readlines() if item.strip()]
        if not stocks_names:
            continue
        try:
            df = pd.read_csv('cal_ops/all.csv')
            stocks = []
            for row in df.index:
                sn = StockNumber(df.loc[row])
                if sn.name not in stocks_names:
                    continue
                stocks.append(sn.ts_code)
        except Exception as e:
            logger.error(e)
            continue
        logger.warning(f't_download {txt_name} {stocks}')
        stocks = download_all.run(stocks, today_date)
        st_queue.put((txt_name, len(stocks), stocks))
    st_queue.put((1, 1, 1))


def t_cal():
    log_dir = os.path.basename(__file__).split('.')[0]
    # 需要个股涨跌图打开，一般只计算一天才打开
    bkd = 200
    start_date = 20301230
    while 1:
        if not st_queue:
            continue
        indus, indus_num, stocks = st_queue.get()
        if (1, 1, 1) == (indus, indus_num, stocks):
            logger.warning('所有计算完成')
            break
        if '持仓' in indus:
            gen_excel = True
        else:
            gen_excel = False
        logger.warning(f't_cal 正在计算【{indus}, {indus_num}】')
        end_date = get_end_date(log_dir, indus, indus_num)
        cal_point1_txt.run(True, gen_excel, log_dir, start_date, end_date, stocks, indus_num, indus=indus, bkd=bkd, N=6)
        logger.warning(f't_cal 计算完成【{indus}, {indus_num}】')


def get_today(d=1):
    if d:
        sz_index_df = ak.index_zh_a_hist(symbol="000001", period="daily")
        df_sorted = sz_index_df.sort_values(by='日期', ascending=False)
        df_sorted.to_csv("shanghai_index.csv", index=False)
    else:
        df_sorted = pd.read_csv('shanghai_index.csv')
    today_date = eval(str(list(df_sorted['日期'])[0]).replace('-', ''))
    return today_date


if __name__ == '__main__':
    try:
        if len(sys.argv) == 2:
            kw = str(sys.argv[1])
        else:
            kw = None
            # kw = '1000持仓'
    except Exception as e:
        logger.error(e)
        kw = '1000持仓'
    logger.add("log/main_all_txt_{time}.log", level='WARNING')
    today_date = get_today()
    t1 = threading.Thread(target=t_download, args=(today_date, kw))
    t2 = threading.Thread(target=t_cal, args=())
    t2.start()
    t1.start()
    t2.join()
    t1.join()
